import logging as log

from sqlalchemy import create_engine, inspect
from jinja2 import Environment, FileSystemLoader
from pathlib import Path

class CodeGenerator:
    def __init__(self, connection_string, template_path, output_base_path):
        self.engine = create_engine(connection_string)
        self.inspector = inspect(self.engine)
        self.template_env = Environment(loader=FileSystemLoader(template_path), trim_blocks=True)
        self.template_path = Path(template_path)  # 使用Path对象
        self.output_base_path = Path(output_base_path)  # 使用Path对象
        self.output_base_path.mkdir(parents=True, exist_ok=True)

    def get_all_tables_info(self):
        tables = []
        for table_name in self.inspector.get_table_names():
            columns = self.inspector.get_columns(table_name)
            tables.append({
                'name': table_name,
                'columns': columns
            })
        return tables

    def render_template(self, template_path, data):
        template = self.template_env.get_template(template_path.name)
        return template.render(data)

    def generate_from_template(self, template_path, data):
        relative_path = template_path.relative_to(self.template_path)
        output_folder = self.output_base_path.joinpath(*relative_path.parts)
        output_folder.mkdir(parents=True, exist_ok=True)
        output_file = output_folder.joinpath(template_path.stem + '.py')
        with output_file.open('w') as file:
            file.write(self.render_template(template_path, data))
        print(f"代码已生成: {output_file}")

    def generate_all(self, data):
        for template_path in self.template_path.glob('**/*.j2'):
            self.generate_from_template(template_path, data)

# 使用示例
if __name__ == "__main__":
    connection_string = '你的数据库连接字符串'
    template_path = Path('templates')  # 使用Path对象
    output_base_path = Path('generated')  # 使用Path对象

    code_gen = CodeGenerator(connection_string, template_path, output_base_path)
    tables_info = code_gen.get_all_tables_info()

    # 生成所有模板文件
    code_gen.generate_all({'tables': tables_info})